# _*_ coding:utf-8 _*_
import pickle
import threading
from abc import ABC, abstractmethod
from random import uniform, randint

import numpy as np
from tqdm import trange


def deepcopy(target):
    return pickle.loads(pickle.dumps(target))


def nearest_neighbor_cost(distance_matrix):
    """ return the cost of route constructing by the nearest neighbor heuristic """
    cost, route = 0, [1]
    unvisited = list(range(2, len(distance_matrix)))
    while unvisited:
        target, min_dis = -1, 1e99
        for city in unvisited:
            if min_dis > distance_matrix[route[-1]][city]:
                min_dis = distance_matrix[route[-1]][city]
                target = city
        cost += min_dis
        route.append(target)
        unvisited.remove(target)
    route.append(route[0])
    cost += distance_matrix[route[-1]][route[0]]
    return cost


class BaseACO(ABC):
    def __init__(self, data):
        self.name = __class__.__name__
        self.data = deepcopy(data)
        self.V = data.cities  # city list
        self.n = len(self.V)  # the number of cities
        self.Distance = data.distance_matrix  # distance matrix
        self.Pheromone = None  # pheromone matrix
        # --- common params, need retuning later ---
        self.alpha = 1.0  # the weight of pheromone intensity
        self.beta = 2.0  # the weight of heuristic info
        self.rho = 0.5  # the rate of pheromone evaporation
        self.m = 0  # the number of ants
        self.tau_ini = None  # the initial pheromone amount
        self.q_zero = 0  # ratio of pseudo random
        self.concurrent = False  # if building routes concurrently
        # --- population ---
        self.colony = []  # the ant colony
        self.optima = self.Ant()  # the best ant ever in history
        self.evolve_list = []  # optimal cost in evolution progress
        self.is_stagnated = False  # if the colony stagnated
        self.stagnation_limit = 100  # iteration limit of stagnation
    
    class Ant:
        def __init__(self):
            self.tour = None
            self.cost = 1e99
    
    def build_route(self, random_start=True):
        """ @return an 'ant' with its tour and traveling cost """
        start = randint(1, self.n) if random_start else 1
        unvisited = deepcopy(self.V)
        unvisited.remove(start)
        ant = self.Ant()
        ant.tour, ant.cost = [start], 0.0
        while unvisited:
            next_city = self.transition(curr=ant.tour[-1], roll=unvisited)
            ant.tour.append(next_city)
            ant.cost += self.Distance[ant.tour[-2]][ant.tour[-1]]
            self.pheromone_update_local(ant.tour[-2], ant.tour[-1])
            unvisited.remove(next_city)
        ant.tour.append(start)
        ant.cost += self.Distance[ant.tour[-2]][ant.tour[-1]]
        return ant
    
    def transition(self, curr, roll):
        """ common transition strategy """
        if len(roll) == 0:
            return 0
        roll = list(roll)
        prob = np.power(self.Pheromone[curr][roll], self.alpha) * np.power(1 / self.Distance[curr][roll], self.beta)
        prob = prob / prob.sum()
        if uniform(0, 1) < self.q_zero:
            city = roll[int(np.argmax(prob))]
        else:
            city = np.random.choice(a=roll, size=1, p=prob)[0]
        return city
    
    class AntThread(threading.Thread):
        def __init__(self, func, args=None):
            threading.Thread.__init__(self)
            self.func = func
            self.args = args
            self.result = None
        
        def run(self) -> None:
            self.result = self.func(self.args)
    
    @abstractmethod
    def pheromone_update_local(self, a, b):
        """ implement this method in subclasses """
    
    @abstractmethod
    def pheromone_update_global(self):
        """ implement this method in subclasses """
    
    def optimize(self, iteration=1000, trial=1):
        stag = 0
        pbar = trange(iteration)
        pbar.set_description(f'{self.__class__.__name__} on {self.data.name} trial {trial}')
        for _ in pbar:
            if self.concurrent:  # build routes concurrently
                threads = [self.AntThread(func=self.build_route) for _ in range(self.m)]
                for thread in threads:
                    thread.start()
                for thread in threads:
                    thread.join()  # waiting for all threads completed
                self.colony = [thread.result for thread in threads]
            else:  # build routes sequentially
                self.colony = [self.build_route() for _ in range(self.m)]
            
            self.colony.sort(key=lambda _: _.cost)
            if self.optima.cost > self.colony[0].cost:
                self.optima = deepcopy(self.colony[0])
                stag = 0
            else:
                stag += 1
            self.is_stagnated = True if stag > self.stagnation_limit else False
            self.evolve_list.append(self.optima.cost)
            self.pheromone_update_global()
            pbar.set_postfix_str(f'best: {self.optima.cost:.2f}')
        return self.optima


class AS(BaseACO):
    def __init__(self, data):
        super(AS, self).__init__(data)
        self.m = self.n
        self.tau_ini = self.m / nearest_neighbor_cost(self.Distance)
        self.Pheromone = np.full(shape=(self.n + 1, self.n + 1), fill_value=self.tau_ini)  # init pheromone matrix
        assert self.Distance.shape == self.Pheromone.shape
    
    def pheromone_update_local(self, a, b):
        pass
    
    def pheromone_update_global(self):
        self.Pheromone = (1 - self.rho) * self.Pheromone  # evaporate
        for ant in self.colony:  # accumulate
            delta = 1.0 / ant.cost
            for i in range(1, len(ant.gene)):
                self.Pheromone[ant.gene[i - 1]][ant.gene[i]] += delta
                self.Pheromone[ant.gene[i]][ant.gene[i - 1]] += delta


class EAS(BaseACO):
    def __init__(self, data):
        super(EAS, self).__init__(data)
        self.m = self.n
        self.e = self.n
        self.tau_ini = (self.e + self.m) / (self.rho * nearest_neighbor_cost(self.Distance))
        self.Pheromone = np.full(shape=(self.n + 1, self.n + 1), fill_value=self.tau_ini)  # init pheromone matrix
        assert self.Distance.shape == self.Pheromone.shape
    
    def pheromone_update_local(self, a, b):
        pass
    
    def pheromone_update_global(self):
        self.Pheromone = (1 - self.rho) * self.Pheromone  # evaporate
        for ant in self.colony:  # accumulate
            delta = 1.0 / ant.cost
            for i in range(1, len(ant.gene)):
                self.Pheromone[ant.gene[i - 1]][ant.gene[i]] += delta
                self.Pheromone[ant.gene[i]][ant.gene[i - 1]] += delta
        delta = self.e / self.optima.cost
        for i in range(1, len(self.optima.tour)):  # accumulate (on the best ever tour additionally)
            self.Pheromone[self.optima.tour[i]][self.optima.tour[i - 1]] += delta
            self.Pheromone[self.optima.tour[i - 1]][self.optima.tour[i]] += delta


class ASRank(BaseACO):
    def __init__(self, data):
        super(ASRank, self).__init__(data)
        self.m = self.n
        self.w = 6
        self.rho = 0.1
        self.tau_ini = 1 / (self.rho * nearest_neighbor_cost(self.Distance))
        self.Pheromone = np.full(shape=(self.n + 1, self.n + 1), fill_value=self.tau_ini)  # init pheromone matrix
        assert self.Distance.shape == self.Pheromone.shape
    
    def pheromone_update_local(self, a, b):
        pass
    
    def pheromone_update_global(self):
        self.Pheromone = (1 - self.rho) * self.Pheromone  # evaporate
        delta = self.w / self.optima.cost
        for i in range(1, len(self.optima.tour)):  # accumulate (on the best ever tour additionally)
            self.Pheromone[self.optima.tour[i]][self.optima.tour[i - 1]] += delta
            self.Pheromone[self.optima.tour[i - 1]][self.optima.tour[i]] += delta
        for r in range(1, self.w):  # accumulate on best ranks
            ant = self.colony[r - 1]
            delta = (self.w - r) / ant.cost
            for i in range(1, len(ant.gene)):
                self.Pheromone[ant.gene[i - 1]][ant.gene[i]] += delta
                self.Pheromone[ant.gene[i]][ant.gene[i - 1]] += delta


class MMAS(BaseACO):
    def __init__(self, data):
        super(MMAS, self).__init__(data)
        self.m = self.n
        self.w = 6
        self.rho = 0.02
        self.avg = self.n / 2.0  # the average number of choices ot ants
        self.tau_max = 1 / (self.rho * nearest_neighbor_cost(self.Distance))
        self.tau_min = self.tau_max * (1 - np.power(0.05, 1 / self.n)) / ((self.avg - 1) * np.power(0.05, 1 / self.n))
        self.tau_ini = self.tau_max
        self.Pheromone = np.full(shape=(self.n + 1, self.n + 1), fill_value=self.tau_ini)  # init pheromone matrix
        assert self.Distance.shape == self.Pheromone.shape
        self.use_current_best = True
    
    def pheromone_update_local(self, a, b):
        pass
    
    def pheromone_update_global(self):
        self.Pheromone = (1 - self.rho) * self.Pheromone  # evaporate
        best = self.colony[0] if self.use_current_best else self.optima
        delta = 1.0 / best.cost
        for i in range(1, len(best.tour)):  # accumulate on 'best' tour
            self.Pheromone[best.tour[i]][best.tour[i - 1]] += delta
            self.Pheromone[best.tour[i - 1]][best.tour[i]] += delta
        if self.is_stagnated:  # stagnation detected, reset pheromone matrix
            self.Pheromone = np.full(shape=(self.n + 1, self.n + 1), fill_value=self.tau_ini)


class ACS(BaseACO):
    def __init__(self, data):
        super(ACS, self).__init__(data)
        self.m = 10
        self.rho = 0.1
        self.xee = 0.1
        self.q_zero = 0.9
        self.tau_ini = 1 / (self.n * nearest_neighbor_cost(self.Distance))
        self.Pheromone = np.full(shape=(self.n + 1, self.n + 1), fill_value=self.tau_ini)  # init pheromone matrix
        assert self.Distance.shape == self.Pheromone.shape
    
    def pheromone_update_local(self, a, b):
        """ update pheromone on edge a-b """
        self.Pheromone[a][b] = (1 - self.xee) * self.Pheromone[a][b]
        self.Pheromone[b][a] = (1 - self.xee) * self.Pheromone[b][a]
        self.Pheromone[a][b] += self.xee * self.tau_ini
        self.Pheromone[b][a] += self.xee * self.tau_ini
    
    def pheromone_update_global(self):
        tour = self.optima.tour  # update only the best ever tour
        delta = 1.0 / self.optima.cost
        for i in range(1, len(tour)):
            self.Pheromone[tour[i - 1]][tour[i]] = (1 - self.rho) * self.Pheromone[tour[i - 1]][tour[i]]
            self.Pheromone[tour[i]][tour[i - 1]] = (1 - self.rho) * self.Pheromone[tour[i]][tour[i - 1]]
            self.Pheromone[tour[i - 1]][tour[i]] += delta
            self.Pheromone[tour[i]][tour[i - 1]] += delta
